'''
convertGraph.py: 
	Take an input graph file and convert it to a suitable format for the
	partitioning/cluster programs (METIS/etc).

-i /Users/cmhill/Work/818K/818k-project/pl_20_2_90 -f METIS -o /Users/cmhill/Downloads/metis-4.0.3/Graphs/pl_20_2_90.metis

$: kmetis /Users/cmhill/Downloads/metis-4.0.3/Graphs/pl_100000_3_90.metis 200
	
Created on Apr 15, 2011
@author: cmhill
'''

import getopt
import networkx as nx
import sys

EDGE = "EDGE"
ADJ_LIST = "ADJ_LIST"

METIS = "METIS"

def createGraph(fileName, format):
	''' Creates and returns a networkx graph object from the input file. '''
	
	graph = nx.Graph()
	
	# Maps nodes
	node_map = {}
	curr_node_counter = 1
	
	file = open(fileName, 'r')
	
	line = file.readline()
	while len(line) is not 0:
		try:
			nodes = line.strip().split('\t')
			
			if format is EDGE:
				# If edge input format:
				node_a, node_b = None, None
				
				if int(nodes[0]) in node_map:
					node_a = node_map[int(nodes[0])]
				else:
					node_a = curr_node_counter
					curr_node_counter += 1
					node_map[int(nodes[0])] = node_a
				
				if int(nodes[1]) in node_map:
					node_b = node_map[int(nodes[1])]
				else:
					node_b = curr_node_counter
					curr_node_counter += 1
					node_map[int(nodes[1])] = node_b
				
				#nodes = line.strip().split('\t')
				
				graph.add_edge(str(node_a), str(node_b))
				#print str(node_a) + ' ' + str(node_b)
				#graph.add_edge(nodes[0], nodes[1])
			else:
				# If adj list format:
				for neighbor in nodes[1:]:
					graph.add_edge(nodes[0], neighbor)
		except:
			print line
			
		line = file.readline()

	return graph

def convertGraph(graph, outputFormat, output):
	''' Convert the networkx graph to the desired output format. '''

	if outputFormat == METIS:
		output.write(str(len(graph.node)) + "\t" + str(len(graph.edges())) + "\n")
		for i in range(1, len(graph.node) + 1):
			#print str(i) + "\t",
			output.write("\t".join(graph.neighbors(str(i))) + "\n")
			
	
	for e in graph.edges():
		print e[0] + '\t' + e[1]


def main():
	try:
		opts, args = getopt.getopt(sys.argv[1:], "af:i:o:v", ["input=", "output=", "adj", "out_format"])
	except getopt.GetoptError, err:
		print str(err) # will print something like "option -a not recognized"
		# usage()
		sys.exit(2)
	
	output = sys.stdout
	outputGraphFormat = None
	inputGraphFormat = EDGE
	inputGraphFile = None
	verbose = False
	
	for o, a in opts:
		if o == "-v":
			verbose = True
		elif o in ("-i", "--input"):
			inputGraphFile = a
		elif o in ("-a", "--adj"):
			inputGraphFormat = ADJ_LIST
		elif o in ("-o", "--output"):
			output = open(a, 'w')
		elif o in ("-f", "--output_format"):
			outputGraphFormat = a
		elif o in ("-h", "--help"):
			# usage()
			sys.exit()
		else:
			assert False, "unhandled option"
	
	# Create the networkx representation of the graph.
	graph = createGraph(inputGraphFile, inputGraphFormat)
	
	# Convert the graph to the desired format.
	convertGraph(graph, outputGraphFormat, output)
	
if __name__ == '__main__':
	main()
